package com.hzgj.bcl.soa.rpc;

import com.google.common.collect.Maps;

import com.hzgj.bcl.soa.Address;
import com.hzgj.bcl.soa.Event;
import com.hzgj.bcl.soa.Interceptor;
import com.hzgj.bcl.soa.rpc.client.Request;
import com.hzgj.bcl.soa.rpc.client.Response;
import com.hzgj.bcl.soa.rpc.message.Message;
import com.hzgj.bcl.soa.rpc.retry.RetryHandler;
import com.hzgj.bcl.soa.rpc.retry.RetryPolicy;
import com.hzgj.bcl.soa.spi.MessageHandleFactory;
import com.hzgj.bcl.soa.util.Constants;
import com.hzgj.bcl.soa.util.ServiceLoaderUtil;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.util.concurrent.ConcurrentMap;


/**
 * Framework to hanle requests and receive matching responses (matching on
 * request ID).
 * Multiple requests can be sent at a time. Whenever a response is received,
 * the correct <code>RspCollector</code> is looked up (key = id) and its
 * method <code>receiveResponse()</code> invoked. A caller may use
 * <code>done()</code> to signal that no more responses are expected, and that
 * the corresponding entry may be removed.
 * <p>
 * <code>RequestCorrelator</code> can be installed at both api and tp
 * sides, it can also switch roles dynamically; i.e., hanle a request and at
 * the same time process an incoming request (when local delivery is enabled,
 * this is actually the default).
 * <p>
 */
public class RequestCorrelator {

    protected static final Logger log = LoggerFactory.getLogger(RequestCorrelator.class);
    /**
     * The table of pending requests (keys=Long (request IDs), values=<tt>RequestEntry</tt>)
     */
    protected final ConcurrentMap<Long, RspCollector> requests = Maps.newConcurrentMap();
    /**
     * The protocol layer to use to pass up/down messages. Can be either a Protocol or a Transport
     */
    protected Interceptor transport;
    /**
     * The handler for the incoming requests. It is called from inside the dispatcher thread
     */
    protected RequestHandler request_handler;
    protected boolean started = false;
    /**
     * Whether or not to use async dispatcher
     */
    protected boolean async_dispatching = false;

    /**
     * Constructor. Uses transport to handle messages. If <code>handler</code>
     * is not null, all incoming requests will be dispatched to it (via
     * <code>handle(HttpMessage)</code>).
     *
     * @param transport Used to hanle/pass up requests. Is a Protocol (up_int.up()/down_int.down() will be used)
     * @param handler   Request handler. method <code>handle(HttpMessage)</code>
     *                  will be called when a request is received.
     */
    public RequestCorrelator(Interceptor transport, RequestHandler handler) {
        this.transport = transport;
        request_handler = handler;
    }

    public void setRequestHandler(RequestHandler handler) {
        request_handler = handler;
    }

    public boolean asyncDispatching() {
        return async_dispatching;
    }

    public RequestCorrelator asyncDispatching(boolean flag) {
        async_dispatching = flag;
        return this;
    }

    /**
     * Sends a request to a single destination
     *
     * @param msg
     * @param coll
     * @throws Exception
     */
    public void sendRequest(long req_id, Message msg, RspCollector coll) throws Exception {
        if (transport == null) {
            if (log.isWarnEnabled()) {
                log.warn("transport is not available !");
            }
            return;
        }
        // i.   Create the request correlator header and add it to the msg
        // ii.  If a reply is expected (coll != null), add a coresponding entry in the pending requests table
        // iii. If deadlock detection is enabled, set/update the call chainStack
        // iv.  Pass the msg down to the protocol layer below

        msg.putHeader(Constants.MESSAGE_SID, req_id);
        if (coll != null) {
            addEntry(req_id, coll);
        }
        transport.down(new Event(Event.MSG, msg));
    }


    /**
     * Used to signal that a certain request may be garbage collected as all responses have been received.
     */
    public void done(long id) {
        removeEntry(id);
    }


    /**
     * <b>Callback</b>.
     * <p>
     * Called by the protocol below when a message has been received. The
     * algorithm should test whether the message is destined for us and,
     * if not, pass it up to the next layer. Otherwise, it should remove
     * the header and check whether the message is a request or response.
     * In the first case, the message will be delivered to the request
     * handler registered (calling its <code>handle()</code> method), in the
     * second case, the corresponding response collector is looked up and
     * the message delivered.
     *
     * @param evt The event to be received
     * @return Whether or not the event was consumed. If true, don't pass message up, else pass it up
     */
    public boolean receive(Event evt) {
        switch (evt.getType()) {
            case Event.MSG:
                if (receiveMessage((Message) evt.getArg())) {
                    return true; // message was consumed, don't pass it up
                }
                break;
            case Event.REINVOKE:
                handleRetry((Message) evt.getArg());
                break;
        }
        return false;
    }

    protected void handleRetry(Message message) {
        RetryPolicy retryPolicy = (RetryPolicy) message.getHeader(Constants.MESSAGE_HEADER_RETRYPOLICY);
        if (retryPolicy == null) {
            Throwable t = (Throwable) message.getHeader(Constants.MESSAGE_HEADER_ATTACHMENT);
            receiveMessage(message.copy(false).setPayload(t).putHeader(Constants.MESSAGE_TYPE, Message.EXC_RSP));
            return;
        }

        Long sid = (Long) message.getHeader(Constants.MESSAGE_SID);
        if (sid == null) {
            log.error("can't find message id of " + message);
        }
        RetryPolicy.InvokeState invokeState = (RetryPolicy.InvokeState) message.getHeader(Constants.MESSAGE_HEADER_RETRY_STATE);
        if (invokeState == null) {
            invokeState = new RetryPolicy.InvokeState() {
            };
        }
        try {
            final RetryPolicy.InvokeState _invokeState = invokeState;
            retryPolicy.retry(invokeState, new RetryHandler() {
                @Override
                public void doRetry() throws IOException {
                    long newId = Request.newId();
                    message.putHeader(Constants.MESSAGE_SID, newId);
                    message.putHeader(Constants.MESSAGE_HEADER_RETRY_STATE, _invokeState);
                    RspCollector rsp = removeEntry(sid);
                    if (rsp == null) {
                        log.error("can't retry,response has return.");
                        return;
                    }
                    addEntry(newId, rsp);
                    log.info("retry handle {} with count {}", message, _invokeState.retryCount);
                    transport.down(new Event(Event.REINVOKE, message));
                }

                @Override
                public void catchException(IOException exception) {
                }
            });
        } catch (RetryPolicy.RetryCancelled retryCancelled) {
            Throwable t = (Throwable) message.getHeader(Constants.MESSAGE_HEADER_ATTACHMENT);
            receiveMessage(message.copy(false).setPayload(t).putHeader(Constants.MESSAGE_TYPE, Message.EXC_RSP));
        }
    }

    public final void start() {
        started = true;
    }

    public void stop() {
        started = false;
        requests.clear();
    }

    /**
     * Handles a message coming from a layer below
     *
     * @return true if the message was consumed, don't pass it further up, else false
     */
    public boolean receiveMessage(Message msg) {
        Long sid = (Long) msg.getHeader(Constants.MESSAGE_SID);
        if (sid == null) {
            throw new RuntimeException("sid can not be null,protocol:" + msg.getHeader(Constants.MESSAGE_HEADER_PROTOCOL));
        }

        // [Header.REQ]:
        // i. If there is no request handler, discard
        // ii. Check whether priority: if synchronous and call chainStack contains
        // address that equals local address -> add priority request. Else
        // add normal request.
        //
        // [Header.RSP]:
        // Remove the msg request correlator header and notify the associated
        // <tt>RspCollector</tt> that a reply has been received
        Integer type = (Integer) msg.getHeader(Constants.MESSAGE_TYPE);
        switch (type) {
            case Message.REQ:
                handleRequest(msg);
                break;
            case Message.RSP:
            case Message.EXC_RSP:
                RspCollector coll = requests.get(sid);
                if (coll != null) {
                    boolean is_exception = type.equals(Message.EXC_RSP);
                    MessageHandleFactory factory = ServiceLoaderUtil.getMessageBuildFactory((String) msg.getHeader(Constants.MESSAGE_HEADER_PROTOCOL));
                    try {
                        coll.receiveResponse(factory.handleResponse(msg.getPayload()), (Address) msg.getHeader(Constants.MESSAGE_ADDRESS_SRC), is_exception);
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
                break;

            default:
                if (log.isErrorEnabled()) {
                    log.error("header's type is neither REQ nor RSP !");
                }
                break;
        }

        return true; // message was consumed
    }

    private void addEntry(long id, RspCollector coll) {
        requests.putIfAbsent(id, coll);
    }


    private RspCollector removeEntry(long id) {
        return requests.remove(id);
    }

    protected void handleRequest(Message req) {
        Object retval;
        boolean threw_exception = false;
        Long sid = (Long) req.getHeader(Constants.MESSAGE_SID);
        if (log.isTraceEnabled()) {
            log.trace(new StringBuilder("calling (").append((request_handler != null ? request_handler.getClass().getName() : "null")).
                    append(") with request ").append(sid).toString());
        }

        if (async_dispatching && request_handler instanceof AsyncRequestHandler) {
            Response rsp = new ResponseImpl(sid);
            try {
                ((AsyncRequestHandler) request_handler).handle(req, rsp);
            } catch (Throwable t) {
                rsp.send(new InvocationTargetException(t), (String) req.getHeader(Constants.MESSAGE_HEADER_PROTOCOL), true);
            }
            return;
        }

        try {
            retval = request_handler.handle(req);
        } catch (Throwable t) {
            log.error("request_handler.handle exception", t);
            threw_exception = true;
            retval = t;
        }
        sendReply(sid, retval, threw_exception, (String) req.getHeader(Constants.MESSAGE_HEADER_PROTOCOL));
    }

    protected void sendReply(final Long req_id, Object reply, boolean is_exception, String protocol) {
        MessageHandleFactory factory = ServiceLoaderUtil.getMessageBuildFactory(protocol);
        Message message = factory.buildResponse(req_id, reply, is_exception);
        message.putHeader(Constants.MESSAGE_TYPE, is_exception ? Message.EXC_RSP : Message.RSP)
                .putHeader(Constants.MESSAGE_SID, req_id)
                .putHeader(Constants.MESSAGE_HEADER_PROTOCOL, protocol);
        sendResponse(message, req_id);
    }

    protected void sendResponse(Message rsp, Long sid) {
        if (log.isDebugEnabled()) {
            log.debug(new StringBuilder("sending rsp for ").append(sid).toString());
        }
        transport.down(new Event(Event.MSG, rsp));
    }

    protected class ResponseImpl implements Response {
        protected final long req_id;

        public ResponseImpl(long req_id) {
            this.req_id = req_id;
        }

        @Override
        public void send(Object reply, String protocol, boolean is_exception) {
            sendReply(req_id, reply, is_exception, protocol);
        }
    }
}
